// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/functions/lib/aggregates/decimal_aggregate.h>
#include <pollux/type/decimal_util.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::functions::aggregate {
    namespace {
        /// Translate selected rows of decoded to the corresponding rows of its base
        /// vector.
        SelectivityVector translateToInnerRows(
            const DecodedVector &decoded,
            const SelectivityVector &rows) {
            POLLUX_DCHECK(!decoded.isIdentityMapping());
            if (decoded.isConstantMapping()) {
                auto constantIndex = decoded.index(rows.begin());
                SelectivityVector baseRows{constantIndex + 1, false};
                baseRows.setValid(constantIndex, true);
                baseRows.updateBounds();
                return baseRows;
            } else {
                SelectivityVector baseRows{decoded.base()->size(), false};
                rows.applyToSelected(
                    [&](auto row) { baseRows.setValid(decoded.index(row), true); });
                baseRows.updateBounds();
                return baseRows;
            }
        }

        /// Return the selected rows of the base vector of decoded corresponding to
        /// rows. If decoded is not identify mapping, baseRowsHolder contains the
        /// selected base rows. Otherwise, baseRowsHolder is unset.
        const SelectivityVector *getBaseRows(
            const DecodedVector &decoded,
            const SelectivityVector &rows,
            SelectivityVector &baseRowsHolder) {
            const SelectivityVector *baseRows = &rows;
            if (!decoded.isIdentityMapping() && rows.hasSelections()) {
                baseRowsHolder = translateToInnerRows(decoded, rows);
                baseRows = &baseRowsHolder;
            }
            return baseRows;
        }

        template<typename TSum>
        struct SumCount {
            TSum sum{0};
            int64_t count{0};
        };
    } // namespace

    /// Partial aggregation produces a pair of sum and count.
    /// Final aggregation produces the average (arithmetic mean)
    /// of all non-null input values.
    template<typename TInput, typename TAccumulator, typename TResult>
    class AverageAggregateBase : public exec::Aggregate {
    public:
        explicit AverageAggregateBase(TypePtr resultType)
            : exec::Aggregate(resultType) {
        }

        int32_t accumulatorFixedWidthSize() const override {
            return sizeof(SumCount<TAccumulator>);
        }

        void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            auto vector = (*result)->as<FlatVector<TResult> >();
            POLLUX_CHECK(vector);
            vector->resize(numGroups);
            uint64_t *rawNulls = getRawNulls(vector);

            TResult *rawValues = vector->mutableRawValues();
            for (int32_t i = 0; i < numGroups; ++i) {
                char *group = groups[i];
                if (isNull(group)) {
                    vector->set_null(i, true);
                } else {
                    clearNull(rawNulls, i);
                    auto *sumCount = accumulator(group);
                    if constexpr (std::is_integral_v<TResult>) {
                        rawValues[i] =
                                static_cast<TResult>(std::round(sumCount->sum / sumCount->count));
                    } else {
                        rawValues[i] = sumCount->sum / sumCount->count;
                    }
                }
            }
        }

        void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            auto row_vector = (*result)->as<RowVector>();
            auto sumVector = row_vector->childAt(0)->as_flat_vector<TAccumulator>();
            auto countVector = row_vector->childAt(1)->as_flat_vector<int64_t>();

            row_vector->resize(numGroups);
            sumVector->resize(numGroups);
            countVector->resize(numGroups);
            uint64_t *rawNulls = getRawNulls(row_vector);

            int64_t *rawCounts = countVector->mutableRawValues();
            TAccumulator *rawSums = sumVector->mutableRawValues();
            for (auto i = 0; i < numGroups; ++i) {
                char *group = groups[i];
                if (isNull(group)) {
                    row_vector->set_null(i, true);
                } else {
                    clearNull(rawNulls, i);
                    auto *sumCount = accumulator(group);
                    rawCounts[i] = sumCount->count;
                    rawSums[i] = sumCount->sum;
                }
            }
        }

        void addRawInput(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            decodedRaw_.decode(*args[0], rows);
            if (decodedRaw_.isConstantMapping()) {
                if (!decodedRaw_.is_null_at(0)) {
                    auto value = decodedRaw_.value_at<TInput>(0);
                    rows.applyToSelected([&](vector_size_t i) {
                        updateNonNullValue(groups[i], TAccumulator(value));
                    });
                }
            } else if (decodedRaw_.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (decodedRaw_.is_null_at(i)) {
                        return;
                    }
                    updateNonNullValue(
                        groups[i], TAccumulator(decodedRaw_.value_at<TInput>(i)));
                });
            } else if (!exec::Aggregate::numNulls_ && decodedRaw_.isIdentityMapping()) {
                auto data = decodedRaw_.data<TInput>();
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue<false>(groups[i], data[i]);
                });
            } else {
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue(
                        groups[i], TAccumulator(decodedRaw_.value_at<TInput>(i)));
                });
            }
        }

        void addSingleGroupRawInput(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            decodedRaw_.decode(*args[0], rows);

            if (decodedRaw_.isConstantMapping()) {
                if (!decodedRaw_.is_null_at(0)) {
                    const TInput value = decodedRaw_.value_at<TInput>(0);
                    const auto numRows = rows.countSelected();
                    updateNonNullValue(group, numRows, TAccumulator(value) * numRows);
                }
            } else if (decodedRaw_.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (!decodedRaw_.is_null_at(i)) {
                        updateNonNullValue(
                            group, TAccumulator(decodedRaw_.value_at<TInput>(i)));
                    }
                });
            } else if (!exec::Aggregate::numNulls_ && decodedRaw_.isIdentityMapping()) {
                const TInput *data = decodedRaw_.data<TInput>();
                TAccumulator totalSum(0);
                rows.applyToSelected([&](vector_size_t i) { totalSum += data[i]; });
                updateNonNullValue<false>(group, rows.countSelected(), totalSum);
            } else {
                TAccumulator totalSum(0);
                rows.applyToSelected(
                    [&](vector_size_t i) { totalSum += decodedRaw_.value_at<TInput>(i); });
                updateNonNullValue(group, rows.countSelected(), totalSum);
            }
        }

        void addIntermediateResults(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /* mayPushdown */) override {
            decodedPartial_.decode(*args[0], rows);
            auto baseRowVector = decodedPartial_.base()->template as<RowVector>();

            if (validateIntermediateInputs_ &&
                (baseRowVector->childAt(0)->may_have_nulls() ||
                 baseRowVector->childAt(1)->may_have_nulls())) {
                addIntermediateResultsImpl<true>(groups, rows);
                return;
            }
            addIntermediateResultsImpl<false>(groups, rows);
        }

        void addSingleGroupIntermediateResults(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /* mayPushdown */) override {
            decodedPartial_.decode(*args[0], rows);
            auto baseRowVector = decodedPartial_.base()->template as<RowVector>();

            if (validateIntermediateInputs_ &&
                (baseRowVector->childAt(0)->may_have_nulls() ||
                 baseRowVector->childAt(1)->may_have_nulls())) {
                addSingleGroupIntermediateResultsImpl<true>(group, rows);
                return;
            }
            addSingleGroupIntermediateResultsImpl<false>(group, rows);
        }

    protected:
        /// Partial.
        template<bool tableHasNulls = true>
        inline void updateNonNullValue(char *group, TAccumulator value) {
            if constexpr (tableHasNulls) {
                exec::Aggregate::clearNull(group);
            }
            accumulator(group)->sum += value;
            accumulator(group)->count =
                    checkedPlus<int64_t>(accumulator(group)->count, 1);
        }

        template<bool tableHasNulls = true>
        inline void updateNonNullValue(char *group, int64_t count, TAccumulator sum) {
            if constexpr (tableHasNulls) {
                exec::Aggregate::clearNull(group);
            }
            accumulator(group)->sum += sum;
            accumulator(group)->count =
                    checkedPlus<int64_t>(accumulator(group)->count, count);
        }

        inline SumCount<TAccumulator> *accumulator(char *group) {
            return exec::Aggregate::value<SumCount<TAccumulator> >(group);
        }

        template<bool checkNullFields>
        void addIntermediateResultsImpl(
            char **groups,
            const SelectivityVector &rows) {
            auto baseRowVector = decodedPartial_.base()->template as<RowVector>();

            SelectivityVector baseRowsHolder;
            auto *baseRows = getBaseRows(decodedPartial_, rows, baseRowsHolder);

            DecodedVector baseSumDecoded{*baseRowVector->childAt(0), *baseRows};
            DecodedVector baseCountDecoded{*baseRowVector->childAt(1), *baseRows};

            if (decodedPartial_.isConstantMapping()) {
                if (!decodedPartial_.is_null_at(0)) {
                    auto decodedIndex = decodedPartial_.index(0);
                    if constexpr (checkNullFields) {
                        POLLUX_USER_CHECK(
                            !baseSumDecoded.is_null_at(decodedIndex) &&
                            !baseCountDecoded.is_null_at(decodedIndex));
                    }
                    auto count = baseCountDecoded.template value_at<int64_t>(decodedIndex);
                    auto sum = baseSumDecoded.template value_at<TAccumulator>(decodedIndex);
                    rows.applyToSelected([&](vector_size_t i) {
                        updateNonNullValue(groups[i], count, sum);
                    });
                }
            } else if (decodedPartial_.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (decodedPartial_.is_null_at(i)) {
                        return;
                    }
                    auto decodedIndex = decodedPartial_.index(i);
                    if constexpr (checkNullFields) {
                        POLLUX_USER_CHECK(
                            !baseSumDecoded.is_null_at(decodedIndex) &&
                            !baseCountDecoded.is_null_at(decodedIndex));
                    }
                    updateNonNullValue(
                        groups[i],
                        baseCountDecoded.template value_at<int64_t>(decodedIndex),
                        baseSumDecoded.template value_at<TAccumulator>(decodedIndex));
                });
            } else {
                rows.applyToSelected([&](vector_size_t i) {
                    auto decodedIndex = decodedPartial_.index(i);
                    if constexpr (checkNullFields) {
                        POLLUX_USER_CHECK(
                            !baseSumDecoded.is_null_at(decodedIndex) &&
                            !baseCountDecoded.is_null_at(decodedIndex));
                    }
                    updateNonNullValue(
                        groups[i],
                        baseCountDecoded.template value_at<int64_t>(decodedIndex),
                        baseSumDecoded.template value_at<TAccumulator>(decodedIndex));
                });
            }
        }

        template<bool checkNullFields>
        void addSingleGroupIntermediateResultsImpl(
            char *group,
            const SelectivityVector &rows) {
            auto baseRowVector = decodedPartial_.base()->template as<RowVector>();

            SelectivityVector baseRowsHolder;
            auto *baseRows = getBaseRows(decodedPartial_, rows, baseRowsHolder);

            DecodedVector baseSumDecoded{*baseRowVector->childAt(0), *baseRows};
            DecodedVector baseCountDecoded{*baseRowVector->childAt(1), *baseRows};

            if (decodedPartial_.isConstantMapping()) {
                if (!decodedPartial_.is_null_at(0)) {
                    auto decodedIndex = decodedPartial_.index(0);
                    if constexpr (checkNullFields) {
                        POLLUX_USER_CHECK(
                            !baseSumDecoded.is_null_at(decodedIndex) &&
                            !baseCountDecoded.is_null_at(decodedIndex));
                    }
                    const auto numRows = rows.countSelected();
                    auto totalCount = checkedMultiply<int64_t>(
                        baseCountDecoded.template value_at<int64_t>(decodedIndex), numRows);
                    auto totalSum =
                            baseSumDecoded.template value_at<TAccumulator>(decodedIndex) *
                            numRows;
                    updateNonNullValue(group, totalCount, totalSum);
                }
            } else if (decodedPartial_.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (!decodedPartial_.is_null_at(i)) {
                        auto decodedIndex = decodedPartial_.index(i);
                        if constexpr (checkNullFields) {
                            POLLUX_USER_CHECK(
                                !baseSumDecoded.is_null_at(decodedIndex) &&
                                !baseCountDecoded.is_null_at(decodedIndex));
                        }
                        updateNonNullValue(
                            group,
                            baseCountDecoded.template value_at<int64_t>(decodedIndex),
                            baseSumDecoded.template value_at<TAccumulator>(decodedIndex));
                    }
                });
            } else {
                TAccumulator totalSum(0);
                int64_t totalCount = 0;
                rows.applyToSelected([&](vector_size_t i) {
                    auto decodedIndex = decodedPartial_.index(i);
                    if constexpr (checkNullFields) {
                        POLLUX_USER_CHECK(
                            !baseSumDecoded.is_null_at(decodedIndex) &&
                            !baseCountDecoded.is_null_at(decodedIndex));
                    }
                    totalCount = checkedPlus<int64_t>(
                        totalCount,
                        baseCountDecoded.template value_at<int64_t>(decodedIndex));
                    totalSum += baseSumDecoded.template value_at<TAccumulator>(decodedIndex);
                });
                updateNonNullValue(group, totalCount, totalSum);
            }
        }

        void initializeNewGroupsInternal(
            char **groups,
            melon::Range<const vector_size_t *> indices) override {
            setAllNulls(groups, indices);
            for (auto i: indices) {
                new(groups[i] + offset_) SumCount<TAccumulator>();
            }
        }

        DecodedVector decodedRaw_;
        DecodedVector decodedPartial_;
    };

    template<typename TUnscaledType>
    class DecimalAverageAggregateBase : public DecimalAggregate<TUnscaledType> {
    public:
        explicit DecimalAverageAggregateBase(TypePtr resultType)
            : DecimalAggregate<TUnscaledType>(resultType) {
        }

        virtual TUnscaledType computeFinalValue(
            functions::aggregate::LongDecimalWithOverflowState *accumulator) final {
            // Handles round-up of fraction results.
            int128_t average{0};
            DecimalUtil::computeAverage(
                average, accumulator->sum, accumulator->count, accumulator->overflow);
            return TUnscaledType(average);
        }
    };

    /// @brief Checks the input type for final aggregation of average.
    /// The input type must be (sum:double/long decimal, count:bigint) struct.
    /// @param type input type for final aggregation of average.
    void checkAvgIntermediateType(const TypePtr &type);
} // namespace kumo::pollux::functions::aggregate
